import pandas as pd
import os
from tushare import pro_api
from datetime import datetime,timedelta
from chinese_calendar import is_workday
from utils import Trade_date
from os import path
from time import sleep,time
from retrying import retry
from concurrent.futures import ThreadPoolExecutor

class Dataset():
    def __init__(self):
        self.ts_token = '9078d7e9d562d1ad22630515a177cf57a9c245fefb09bd62f5b07b20'
    
    #@retry(stop_max_attempt_number=5,wait_fixed=60) # 下载失败尝试5次，等待60s
    def download_by_date(self,days=100,start_day=None,end_day=None):
        """下载日线数据,从tushare获取
        
        Args:
            如果结束日期end_day为空，则结束日期默认为今天
            如果开始日期start_day为空，则返回end_day前days个交易日的数据
            如果开始日期start_day不为空，则返回start_day到end_day的交易日数据
            start_day和end_day为字符串，格式为YYYYMMDD
            days类似为整数
        """
        # 获取上一个交易日的迭代器
        last_day = Trade_date(end_day).last_trade_day()
        # 调用tushare 接口获取数据
        pro = pro_api(self.ts_token)
        # 保存到data/下
        if not path.exists('ts_data'):
            os.mkdir('ts_data')
        # 如果开始时间不为空，则按日期下载
        if start_day:
            # 输入的日期格式是否正确
            assert type(start_day) == str and start_day.isdigit() and len(start_day) == 8,'日期格式错误'
            while(True):
                # 获取上一个交易日
                trade_day = next(last_day)
                # 小于开始日期则退出
                if trade_day < start_day:
                    break
                # 下载数据
                file_path = f'ts_data/{trade_day}.csv'
                if not path.exists(file_path):
                    df1 = pro.daily(trade_date=trade_day)
                    df1.to_csv(file_path,index=False)
                print(f'\r正在下载 {trade_day} 的日线数据',end='',flush=True)
                
        else:# 如果开始时间为空，则按天数下载
            for i in range(days):
                # 获取上一个交易日
                trade_day = next(last_day)
                # 下载数据
                file_path = f'ts_data/{trade_day}.csv'
                if not path.exists(file_path):
                    df1 = pro.daily(trade_date=trade_day)
                    df1.to_csv(file_path,index=False)
                print('\r正在下载日线数据,已下载:{}天'.format(i+1),end='',flush=True)
    
    
    def load_data(self,file_path):
        df = pd.read_csv(file_path)
        return df
        
    def get_by_code(self,code,days=1,start_day=None,end_day=None):
        """
        如果结束日期end_day为空，则结束日期默认为今天
        如果开始日期start_day为空，则返回end_day前days个交易日的数据
        如果开始日期start_day不为空，则返回start_day到end_day的交易日数据
        code为股票代码，类似为字符串
        start_day和end_day为字符串，格式为YYYYMMDD
        days类似为整数
        """        
        # 获取上一个交易日的迭代器
        last_day = Trade_date(end_day).last_trade_day()
        # 保存要查询的数据
        df = pd.DataFrame()
        # 如果开始时间不为空，则按日期
        if start_day:
            # 输入的日期格式是否正确
            assert type(start_day) == str and start_day.isdigit() and len(start_day) == 8,'日期格式错误'
            # 太慢了，改用多线程
            tasks = []
            def one_task(trade_day):
                # 查询本地是否有该日数据,没有就先下载下来
                file_path = f'ts_data/{trade_day}.csv'
                if not path.exists(file_path):
                    self.download_by_date(start_day=trade_day,end_day=trade_day)
                # 读取数据，取出该股的数据，然后拼接
                df1 = pd.read_csv(file_path,encoding='utf-8')
                df1 = df1.loc[df1['ts_code']==code]
                return df1
            with ThreadPoolExecutor(12) as expool:
                while(True):
                    # 获取上一个交易日
                    trade_day = next(last_day)
                    # 小于开始日期则退出
                    if trade_day < start_day:
                        break
                    tasks.append(expool.submit(one_task,trade_day))
                for task in tasks:
                    df1 = task.result()
                    df = pd.concat([df,df1])
                
        else:# 如果开始时间为空，则按天数
            # 太慢了，改用多线程
            tasks = []
            def one_task():
                global df
                # 获取上一个交易日
                trade_day = next(last_day)
                # 查询本地是否有该日数据,没有就先下载下来
                file_path = f'ts_data/{trade_day}.csv'
                if not path.exists(file_path):
                    self.download_by_date(start_day=trade_day,end_day=trade_day)
                # 读取数据，取出该股的数据，然后拼接
                df1 = pd.read_csv(file_path,encoding='utf-8')
                df1 = df1.loc[df1['ts_code']==code]
                return df1
            with ThreadPoolExecutor(12) as expool:
                for i in range(days):
                    tasks.append(expool.submit(one_task))
                for task in tasks:
                    df1 = task.result()
                    df = pd.concat([df,df1])
        return df.reset_index(drop=True)
    
    def get_day_line(self,df:pd.DataFrame,days=5):
        """
        获取days日均线，默认获取5日均线
        df为要计算均线的日线数据，类型：DataFrame
        均线长度视传入的df数据量而定
        """
        # 股票种类查重和按日期去重
        assert len(df.groupby('ts_code').count()) == 1,'df中只能包含1只股票数据'
        df = df.drop_duplicates(subset='trade_date')
        # 按日期升序排序，计算均线并丢弃na值，默认前两列为股票代码和日期，后面列为数据
        df = df.sort_values(by='trade_date',ascending=True)
        df = pd.concat([df.iloc[:,:2],df.iloc[:,2:].rolling(5).mean()],axis=1)
        return df.dropna().reset_index(drop=True)
        
    def get_stock_basic(self,**kargs):
        """
        参数列表，详见tushare
        is_hs	    str	N	是否沪深港通标的，N否 H沪股通 S深股通
        list_status	str	N	上市状态 L上市 D退市 P暂停上市，默认是L
        exchange	str	N	交易所 SSE上交所 SZSE深交所 BSE北交所
        ts_code	    str	N	TS股票代码
        market	    str	N	市场类别 （主板/创业板/科创板/CDR/北交所）
        limit	    int	N	
        offset	    int	N	
        name	    str	N	名称
        """
        pro = pro_api(self.ts_token)
        return pro.stock_basic(**kargs)
    
    def get_by_code_date(self,ts_code,trade_date):
        # 按照日期和股票代码获取当天数据
        assert Trade_date().is_trade_day(trade_date),f'{trade_date}不是交易日'
        file_path = f'ts_data/{trade_date}.csv'
        if not path.exists(file_path):
            self.download_by_date(1,trade_date,trade_date)
        df = pd.read_csv(file_path)
        return df[df['ts_code']==ts_code]
            
        

if __name__ == '__main__':
    Dataset().download_by_date()
